(* mlton_ant.sml     *)
(* Eric Rollins 2006 *)

(* 
   Written for the MLton version of Standard ML, http://mlton.org

   This program generates a random array of distances between cities, then uses
   Ant Colony Optimization to find a short path traversing all the cities -- 
   the Travelling Salesman Problem.

   In this version of Ant Colony Optimization each ant starts in a random city.
   Paths are randomly chosed with probability inversely proportional to to the
   distance to the next city.  At the end of its travel the ant updates the 
   pheromone matrix with its path if this path is the shortest one yet found.
   The probability of later ants taking a path is increased by the pheromone 
   value on that path.  Pheromone values evaporate (decrease) over time.

   In this impementation weights between cities actually represent 
   (maxDistance - dist), so we are trying to maximize the score.

   Usage: ant seed boost iterations cities
     seed         seed for random number generator (1,2,3...).
                  This seed controls the city distance array.  Remote 
                  executions have their seed values fixed (1,2) so each will 
                  produce a different result.
     boost        pheromone boost for best path.  5 appears good.  
                  0 disables pheromones, providing random search.
     iterations   number of ants to be run.
     cities       number of cities.
*)

open Array;
open Array2;

structure Main =
struct
   val (seedStr, boostStr, iterStr, numCitiesStr) =
	case CommandLine.arguments ()
	 of [s1, s2, s3, s4] => (s1, s2, s3, s4)
	  | _ => (TextIO.output (TextIO.stdErr, 
                    "usage: ant seed boost iterations cities\n");
		  OS.Process.exit OS.Process.failure)

  val seed = Option.getOpt(Int.fromString(seedStr),1);
  val boost = Option.getOpt(Int.fromString(boostStr),5);
  val iter = Option.getOpt(Int.fromString(iterStr),20);
  val numCities = Option.getOpt(Int.fromString(numCitiesStr),22);

  (* Rand was written for Alice SML; reused here                       *)
  (* Values taken from http://www.math.utah.edu/~pa/Random/Random.html *)
  (* But Alice int is 31 bits, so have to mod again.                   *)
  fun myRandLimits() =
    (0, Option.getOpt(Int.maxInt,32767))

  (* Returns a function which, when called, returns next random number. *)
  (* Current value (initially seedX) is stored in closure.              *)
  fun makeRandomGenerator(seedX) = 
    let
      val x = ref(Int.toLarge(seedX));
      val p1 = Int.toLarge(16807);
      val n  = Int.toLarge( 147483647) +
               Int.toLarge(1000000000) +
               Int.toLarge(1000000000); 
      val n2 = Int.toLarge(#2(myRandLimits()));
      fun f() =
        let 
          val nextX = #2(IntInf.divMod((!x * p1), n));
        in
          x := nextX;
          Int.fromLarge(#2(IntInf.divMod(nextX, n2)))
        end
    in
      f
    end
      
  (* generates random int between 0 and upperBound *)
  fun intRandBound(upperBound, rgen) =
    floor(real(upperBound) * (real(rgen()) / 
      (real(#2(myRandLimits())) + 1.0)));

  fun randBound(upperBound, rgen) =
    upperBound * (real(rgen()) / (real(#2(myRandLimits())) + 1.0));

  (* A simple implementation of fixed-size sets using Array *)
  fun set(n) = Array.array(n,false)
  fun add(set, item) = Array.update(set,item,true)
  fun has(set, item) = Array.sub(set,item)
  fun count(set) = 
    let
      fun f(a,b) = if a then (b+1) else b
    in
      Array.foldl f 0 set
    end
  
  (* matrix is square two-dimensional array *)
  fun printMatrix(arr) =
    let
      fun f(r, c, v) = (
          print(Int.toString(r));
          print(",");
          print(Int.toString(c));
          print(" : ");
          print(Real.toString(v));
          print("\n")
      )
    in
      appi RowMajor f {base=arr,row=0,col=0,nrows=NONE,ncols=NONE}
    end

  (* upperBound must be > 1.0 *)
  fun randomizeMatrix(arr, upperBound, rgen) =
    let
      fun f(v) = randBound(upperBound - 1.0, rgen) + 1.0
    in
      modify RowMajor f arr
    end

  fun printList(lst) =
    let
      fun f(x) = (
        print " ";
        print(Int.toString(x))
      )
    in
      List.app f lst
    end

  fun sumList(lst) = List.foldl Real.+ 0.0 lst

  fun pathLength(cities, path) =
    let
      fun f(r,c) = sub(cities, r, c)
    in
      sumList(ListPair.mapEq f (path, tl(path)@[hd(path)]))
    end 

  fun genPathRecurse(cities, pher, used, path, current, hits, rgen) =
    let
      val c = ref 0;
      val next = ref 0;
      val sumWeight = ref 0.0;
      val rndValue = ref 0.0;
    in
      if count(used) = nCols(cities) then
        (path, hits)
      else (
        while !c < nCols(cities) do (
          if has(used, !c) then
            () (* NOP *)
          else
            sumWeight := 
              !sumWeight + sub(cities, current, !c) * 
	        (1.0 + sub(pher, current, !c));
          c := !c + 1
        );
        rndValue := randBound(!sumWeight, rgen);
        c := 0;
        sumWeight := 0.0;
        while ((!c < nCols(cities)) andalso 
            (has(used, !c) orelse (!sumWeight < !rndValue))) do (
          if has(used, !c) then
            () (* NOP *)
          else (
            sumWeight := 
              !sumWeight + sub(cities, current, !c)  * 
	        (1.0 + sub(pher, current, !c));
            next := !c
          );
          c := !c + 1
        );
        add(used, !next);
        let
          val newHits = 
	    if sub(pher, current, !next) > 0.0 then hits + 1 else hits
        in
          genPathRecurse(cities, pher, used, path@[!next], !next, newHits, 
            rgen)
        end
      )  
    end 

  exception MyNotFound;

  fun genPath(cities, pher, rgen) =
    let  
      val used = set(numCities);
      val start : int = intRandBound(nCols(cities), rgen);
    in
      add(used, start);
      genPathRecurse(cities, pher, used, [start], start, 0, rgen)
    end

  fun updatePher(pher, path) = 
    let
      val incr = real(boost);
      fun f(r,c) = update(pher, r, c, sub(pher, r, c) + incr)
    in
      ListPair.mapEq f (path, tl(path)@[hd(path)])
    end 

  fun evaporatePher(pher) = 
    let
      val decr = real(boost)/real(iter);
      fun f(v) = 
        if v >= decr then
          v - decr
        else
          0.0
    in
      modify RowMajor f pher
    end

  fun work() =
    let
      val maxDistance = 100.0;
      (* Distances between cities (actually maxDistance - dist) *)
      (* cities[r,c] == distance from r to c *)
      (* Randomized values are not symmetric (A->B != B->A) and don't satisfy
         triangle inequality *)
      val cities = array(numCities, numCities, maxDistance);
      val localRgen = makeRandomGenerator(seed);
    in
      randomizeMatrix(cities, maxDistance, localRgen);
      let 
        (* this will be executed on remote host *)
        fun remoteWork(remoteSeed) =            
          let 
            (* Pheromone values *)
            (* Desirability of route from r to c *)
            val pher = array(numCities, numCities, 0.0);
            val i = ref 0;
            val bestLen = ref 0.0;
            val bestPath = ref [];
            val rgen = makeRandomGenerator(remoteSeed);
          in
            while !i < iter do
              let
                val t = genPath(cities, pher, rgen);
                val path = #1(t);
                val hits = #2(t);
                val len = pathLength(cities, path);
              in
(*
                printList(path);
                print " : ";
                print(Real.toString(len));
                print " (";
                print (Int.toString(hits));
                print ")\n";
*)
                if len > !bestLen then (
                  updatePher(pher, path);
                  bestLen := len;
                  bestPath := path
                ) else
                  () (* NOP *);
                evaporatePher(pher);
                i := !i + 1
              end;
            !bestPath
          end
      in
          let
              val bestPath = remoteWork(2);
              val bestLen = pathLength(cities, bestPath)
          in
              printList(bestPath);
              print " : ";
              print(Real.toString(bestLen));
              print "\n"
          end
      end
    end

  val _ = (
             let
               val t = Timer.startRealTimer();
             in
               work(); 
               print(Time.toString(Timer.checkRealTimer(t)));
               print(" seconds\n");
               OS.Process.exit(OS.Process.success)
             end
           )
end